import pandas as pd
import torch
from torch.nn.utils.rnn import pad_sequence
from config import Config


def load_data():
    # 加载原始数据
    df = pd.read_csv(
        '../../tcdata/oppo_breeno_round1_data/train.tsv',
        sep='\t',
        names=['query1', 'query2', 'label'],
        dtype={'query1': str, 'query2': str}
    )

    # 将数字字符串转为ID列表
    df['query1_ids'] = df['query1'].apply(lambda x: list(map(int, x.split())))
    df['query2_ids'] = df['query2'].apply(lambda x: list(map(int, x.split())))

    return df


class QueryPairDataset(torch.utils.data.Dataset):
    def __init__(self, df):
        # 过滤超出VOCAB_SIZE的token
        def filter_ids(ids):
            return [i for i in ids if i < Config.VOCAB_SIZE]

        self.query1 = [torch.tensor(filter_ids(ids)) for ids in df['query1_ids']]
        self.query2 = [torch.tensor(filter_ids(ids)) for ids in df['query2_ids']]
        self.labels = df['label'].values
        self.length = len(df)

    def __len__(self):
        return self.length  # 返回数据集长度

    def __getitem__(self, idx):
        return {
            'query1_ids': self.query1[idx],
            'query2_ids': self.query2[idx],
            'labels': torch.tensor(self.labels[idx], dtype=torch.long)
        }

    @staticmethod
    def collate_fn(batch):
        q1 = pad_sequence(
            [x['query1_ids'].long() for x in batch],  # 添加.long()确保类型
            batch_first=True
        )
        q2 = pad_sequence(
            [x['query2_ids'].long() for x in batch],  # 添加.long()确保类型
            batch_first=True
        )
        return {
            'query1_ids': q1,
            'query2_ids': q2,
            'labels': torch.stack([x['labels'] for x in batch])
        }